import logging

from utils import getenv, set_env_variables
from fastapi import FastAPI, Request, status, HTTPException, Depends
from web.main import app as webui_app
from ollama.main import app as ollama_app
from audio.main import app as audio_app
from images.main import app as images_app
from fastapi.staticfiles import StaticFiles
from rag.main import app as rag_app
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.middleware.base import BaseHTTPMiddleware
import json
import sys
from rag.utils import rag_messages
from expiry import check_expiry


from config import (
    WEBUI_NAME,
    CONFIG_DATA,
    VERSION, SRC_LOG_LEVELS,
    CACHE_DIR,
    FRONTEND_BUILD_DIR,
    GLOBAL_LOG_LEVEL,
    ERROR_MESSAGES

)


logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MAIN"])


class SPAStaticFiles(StaticFiles):
    async def get_response(self, path: str, scope):
        try:
            return await super().get_response(path, scope)
        except (HTTPException, StarletteHTTPException) as ex:
            if ex.status_code == 404:
                return await super().get_response("index.html", scope)
            else:
                raise ex


app = FastAPI()


class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # if "/expiry" not in request.url.path:

            # return await call_next(request)
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
            log.debug(f"request.url.path: {request.url.path}")

            expiry_info = check_expiry()
            if expiry_info.get('trial_remain') == 0:
                raise HTTPException(
                     status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=ERROR_MESSAGES.TRIAL_EXPIRY
                    )
            # Read the original request body
            body = await request.body()
            # Decode body to string
            body_str = body.decode("utf-8")
            # Parse string to JSON
            data = json.loads(body_str) if body_str else {}

            # Example: Add a new key-value pair or modify existing ones
            # data["modified"] = True  # Example modification
            if "docs" in data:
                data = {**data}
                data["messages"] = rag_messages(
                    data["docs"],
                    data["messages"],
                    rag_app.state.RAG_TEMPLATE,
                    rag_app.state.TOP_K,
                    rag_app.state.RAG_EMBEDDING_ENGINE,
                    rag_app.state.RAG_EMBEDDING_MODEL,
                    rag_app.state.sentence_transformer_ef,
                    # rag_app.state.RAG_OPENAI_API_KEY,
                    # rag_app.state.RAG_OPENAI_API_BASE_URL,
                    "",
                    "",
                )
                del data["docs"]

                log.debug(f"data['messages']: {data['messages']}")

            modified_body_bytes = json.dumps(data).encode("utf-8")

            # Replace the request body with the modified one
            request._body = modified_body_bytes

            # Set custom header to ensure content-length matches new body length
            request.headers.__dict__["_list"] = [
                (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
                *[
                    (k, v)
                    for k, v in request.headers.raw
                    if k.lower() != b"content-length"
                ],
            ]

        response = await call_next(request)
        return response

    async def _receive(self, body: bytes):
        return {"type": "http.request", "body": body, "more_body": False}


app.add_middleware(RAGMiddleware)


app.mount("/api/v1", webui_app)
app.mount("/ollama", ollama_app)
app.mount("/images/api/v1", images_app)
app.mount("/audio/api/v1", audio_app)
app.mount("/rag/api/v1", rag_app)

@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


@app.get("/api/config")
async def get_app_config():
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA

    default_locale = "zh-CN"
    if "ui" in CONFIG_DATA:
        default_locale = CONFIG_DATA["ui"].get("default_locale", "zh-CN")

    # The Rest of the Function Now Uses the Variables Defined Above
    return {
        "status": True,
        "name": WEBUI_NAME,
        "version": VERSION,
        "default_locale": default_locale,
        "images": images_app.state.check_image_enable(),
        "default_models": webui_app.state.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
        "trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
    }


@app.get("/expiry")
async def expiry():
    return check_expiry()

app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")

app.mount(
    "/",
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
    name="spa-static-files",
)

if __name__ == "__main__":
    import uvicorn
    import webbrowser
    web_port = int(getenv("PORT", 10004))
    uvicorn.run(app, host="0.0.0.0", port=web_port)
    webbrowser.open("http://localhost:" + str(web_port))